// Copyright (c) 2020, Qihoo, Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef TENSORNET_PS_TABLE_DENSE_TABLE_H_
#define TENSORNET_PS_TABLE_DENSE_TABLE_H_

#include "core/ps/optimizer/optimizer.h"

#include <map>
#include <vector>
#include <mutex>

namespace butil {
class IOBuf;
} // namespace butil

namespace tensornet {

class DenseTable
{
public:
    DenseTable(const OptimizerBase* opt, int shard_num, int self_shard_id);

    ~DenseTable() = default;

    int Init(int total_element);

    int SetWeight(butil::IOBuf& w_buf);

    void SetHandle(uint32_t handle);

    uint32_t GetHandle() const {
        return handle_;
    }

    int TotalElements() {
        return total_elements_;
    }

    const DenseOptKernelSharedPtr GetOptKernels(int shard_id) const;

    void Save(std::string filepath) const;

    void Load(std::string filepath);

private:
    int shard_num_ = 0;
    int self_shard_id_ = 0;
    uint32_t handle_ = 0;
    int total_elements_ = 0;
    const OptimizerBase* opt_ = nullptr;

    std::vector<DenseOptKernelSharedPtr> opt_kernels_;

    bool is_initialized_ = false;
};

class DenseTableRegistry {
public:
    ~DenseTableRegistry() {
        for (auto table : tables_) {
            delete table;
        }
    }

    static DenseTableRegistry* Instance();

    uint32_t Register(DenseTable* table);

    DenseTable* Get(uint32_t table_handle);

protected:
    DenseTableRegistry() {}

private:
    std::mutex table_mu_;
    std::vector<DenseTable*> tables_;
};

DenseTable* CreateDenseTable(const OptimizerBase* opt, int shard_num, int self_shard_id);

} // namespace tensornet {

#endif // !TENSORNET_PS_TABLE_DENSE_TABLE_H_
